/* Copyright 2007-2008 dnAnalytics Project.
 *
 * Contributors to this file:
 * Patrick van der Velde
 *
 * This file is part of dnAnalytics.  dnAnalytics is licensed under the 
 * Microsoft Public License. See License.txt for a complete copy of the
 * license.
 */

using System;
using System.Collections.Generic;
using dnAnalytics.LinearAlgebra.Solvers.Preconditioners;
using dnAnalytics.Math;

namespace dnAnalytics.LinearAlgebra.Solvers.Iterative
{
    /// <summary>
    /// A Generalized Product Bi-Conjugate Gradient iterative matrix solver.
    /// </summary>
    /// <remarks>
    /// <para>
    /// The Generalized Product Bi-Conjugate Gradient (GPBiCG) solver is an 
    /// alternative version of the Bi-Conjugate Gradient stabilized (CG) solver.
    /// Unlike the CG solver the GPBiCG solver can be used on 
    /// non-symmetric matrices. <br/>
    /// Note that much of the success of the solver depends on the selection of the
    /// proper preconditioner.
    /// </para>
    /// <para>
    /// The GPBiCG algorithm was taken from: <br/>
    /// GPBiCG(m,l): A hybrid of BiCGSTAB and GPBiCG methods with 
    /// efficiency and robustness
    /// <br/>
    /// S. Fujino
    /// <br/>
    /// Applied Numerical Mathematics, Volume 41, 2002, pp 107 - 117
    /// <br/>
    /// </para>
    /// <para>
    /// The example code below provides an indication of the possible use of the
    /// solver.
    /// </para>
    /// </remarks>
    /// <include file='../../../../examples/LinearAlgebra/Solvers/Iterative/GPBiCG.xml' path='example'/> 
    public sealed class GPBiCG : IIterativeSolver
    {
        /// <summary>
        /// The status used if there is no status, i.e. the solver hasn't run yet and there is no
        /// iterator.
        /// </summary>
        private static readonly ICalculationStatus sm_DefaultStatus = new CalculationIndetermined();

        /// <summary>
        /// The preconditioner that will be used. Can be set to <c>null</c>, in which case the default
        /// pre-conditioner will be used.
        /// </summary>
        private IPreConditioner m_Preconditioner;

        /// <summary>
        /// The iterative process controller.
        /// </summary>
        private IIterator m_Iterator;

        /// <summary>
        /// Indicates the number of <c>BiCGStab</c> steps should be taken 
        /// before switching.
        /// </summary>
        private int m_NumberOfBiCgStabSteps = 1;

        /// <summary>
        /// Indicates the number of <c>GPBiCG</c> steps should be taken 
        /// before switching.
        /// </summary>
        private int m_NumberOfGPBiCGSteps = 4;

        /// <summary>
        /// Indicates if the user has stopped the solver.
        /// </summary>
        private bool m_HasBeenStopped;

        /// <summary>
        /// Initializes a new instance of the <c>GPBiCG</c> class.
        /// </summary>
        /// <remarks>
        /// When using this constructor the solver will use the <see cref="IIterator"/> with
        /// the standard settings and a default preconditioner.
        /// </remarks>
        public GPBiCG() : this(null, null)
        {}

        /// <summary>
        /// Initializes a new instance of the <c>GPBiCG</c> class.
        /// </summary>
        /// <remarks>
        /// <para>
        /// When using this constructor the solver will use a default preconditioner.
        /// </para>
        /// <para>
        /// The main advantages of using a user defined <see cref="IIterator"/> are:
        /// <list type="number">
        /// <item>It is possible to set the desired convergence limits.</item>
        /// <item>
        /// It is possible to check the reason for which the solver finished 
        /// the iterative procedure by calling the <see cref="IIterator.Status"/> property.
        /// </item>
        /// </list>
        /// </para>
        /// </remarks>
        /// <param name="iterator">
        ///   The <see cref="IIterator"/> that will be used to monitor the iterative process.
        /// </param>
        public GPBiCG(IIterator iterator) : this(null, iterator)
        {}

        /// <summary>
        /// Initializes a new instance of the <c>GPBiCG</c> class.
        /// </summary>
        /// <remarks>
        /// When using this constructor the solver will use the <see cref="IIterator"/> with
        /// the standard settings.
        /// </remarks>
        /// <param name="preconditioner">
        /// The <see cref="IPreConditioner"/> that will be used to precondition the matrix equation.
        /// </param>
        public GPBiCG(IPreConditioner preconditioner) : this(preconditioner, null)
        {}

        /// <summary>
        /// Initializes a new instance of the <c>GPBiCG</c> class.
        /// </summary>
        /// <remarks>
        /// <para>
        /// The main advantages of using a user defined <see cref="IIterator"/> are:
        /// <list type="number">
        /// <item>It is possible to set the desired convergence limits.</item>
        /// <item>
        /// It is possible to check the reason for which the solver finished 
        /// the iterative procedure by calling the <see cref="IIterator.Status"/> property.
        /// </item>
        /// </list>
        /// </para>
        /// </remarks>
        /// <param name="preconditioner">
        /// The <see cref="IPreConditioner"/> that will be used to precondition the matrix equation.
        /// </param>
        /// <param name="iterator">
        /// The <see cref="IIterator"/> that will be used to monitor the iterative process.
        /// </param>
        public GPBiCG(IPreConditioner preconditioner, IIterator iterator)
        {
            m_Iterator = iterator;
            m_Preconditioner = preconditioner;
        }

        /// <summary>
        /// Gets or sets the number of steps taken with the <c>BiCgStab</c> algorithm
        /// before switching over to the <c>GPBiCG</c> algorithm.
        /// </summary>
        public int NumberOfBiCgStabSteps
        {
            get
            {
                return m_NumberOfBiCgStabSteps;
            }
            set
            {
                if (value < 0)
                {
                    throw new ArgumentOutOfRangeException("value");
                }

                m_NumberOfBiCgStabSteps = value;     
            }
        }

        /// <summary>
        /// Gets or sets the number of steps taken with the <c>GPBiCG</c> algorithm
        /// before switching over to the <c>BiCgStab</c> algorithm.
        /// </summary>
        public int NumberOfGpBiCgSteps
        {
            get
            {
                return m_NumberOfGPBiCGSteps;
            }
            set
            {
                if (value < 0)
                {
                    throw new ArgumentOutOfRangeException("value");
                }

                m_NumberOfGPBiCGSteps = value;
            }
        }

        /// <summary>
        /// Sets the <c>IPreConditioner</c> that will be used to precondition the iterative process.
        /// </summary>
        /// <param name="preconditioner">The preconditioner.</param>
        public void SetPreconditioner(IPreConditioner preconditioner)
        {
            m_Preconditioner = preconditioner;
        }

        /// <summary>
        /// Sets the <c>IIterator</c> that will be used to track the iterative process.
        /// </summary>
        /// <param name="iterator">The iterator.</param>
        public void SetIterator(IIterator iterator)
        {
            m_Iterator = iterator;
        }

        /// <summary>
        /// Gets the status of the iteration once the calculation is finished.
        /// </summary>
        public ICalculationStatus IterationResult
        {
            get 
            {
                return (m_Iterator != null) ? m_Iterator.Status : sm_DefaultStatus;
            }
        }

        /// <summary>
        /// Stops the solve process. 
        /// </summary>
        /// <remarks>
        /// Note that it may take an indetermined amount of time for the solver to actually
        /// stop the process.
        /// </remarks>
        public void StopSolve()
        {
            m_HasBeenStopped = true;
        }

        /// <summary>
        /// Solves the matrix equation Ax = b, where A is the coefficient matrix, b is the
        /// solution vector and x is the unknown vector.
        /// </summary>
        /// <param name="matrix">The coefficient matrix, <c>A</c>.</param>
        /// <param name="vector">The solution vector, <c>b</c>.</param>
        /// <returns>The result vector, <c>x</c>.</returns>
        public Vector Solve(Matrix matrix, Vector vector)
        {
            if (vector == null)
            {
                throw new ArgumentNullException();
            }

            Vector result = new DenseVector(matrix.Rows);
            Solve(matrix, vector, result);
            return result;
        }

        /// <summary>
        /// Solves the matrix equation Ax = b, where A is the coefficient matrix, b is the
        /// solution vector and x is the unknown vector.
        /// </summary>
        /// <param name="matrix">The coefficient matrix, <c>A</c>.</param>
        /// <param name="input">The solution vector, <c>b</c></param>
        /// <param name="result">The result vector, <c>x</c></param>
        public void Solve(Matrix matrix, Vector input, Vector result)
        {
            // If we were stopped before, we are no longer
            // We're doing this at the start of the method to ensure
            // that we can use these fields immediately.
            {
                m_HasBeenStopped = false;
            }

            // Error checks
            {
                if (matrix == null)
                {
                    throw new ArgumentNullException("matrix");
                }

                if (matrix.Rows != matrix.Columns)
                {
                    throw new MatrixNotSquareException();
                }

                if (input == null)
                {
                    throw new ArgumentNullException("input");
                }

                if (result == null)
                {
                    throw new ArgumentNullException("result");
                }

                if (result.Count != input.Count)
                {
                    throw new NotConformableException();
                }
            }

            // Initialize the solver fields
            {
                // Set the convergence monitor
                if (m_Iterator == null)
                {
                    m_Iterator = Iterator.CreateDefault();
                }

                if (m_Preconditioner == null)
                {
                    m_Preconditioner = new UnitPreconditioner();
                }
                m_Preconditioner.Initialize(matrix);
            }

            // x_0 is initial guess
            // Take x_0 = 0
            Vector xTemp = new DenseVector(input.Count);

            // r_0 = b - Ax_0
            // This is basically a SAXPY so it could be made a lot faster
            Vector residuals = new DenseVector(matrix.Rows);
            CalculateTrueResidual(matrix, residuals, xTemp, input);

            // Define the temporary scalars
            double alpha = 0;
            double beta = 0;
            double sigma = 0;
            double eta = 0;

            // Define the temporary vectors
            // rDash_0 = r_0
            Vector rDash = new DenseVector(residuals);

            // t_-1 = 0
            Vector t = new DenseVector(residuals.Count);
            Vector t0 = new DenseVector(residuals.Count);
            
            // w_-1 = 0
            Vector w = new DenseVector(residuals.Count);
            
            // Define the remaining temporary vectors
            Vector c = new DenseVector(residuals.Count);
            Vector p = new DenseVector(residuals.Count);
            Vector s = new DenseVector(residuals.Count);
            Vector u = new DenseVector(residuals.Count);
            Vector y = new DenseVector(residuals.Count);
            Vector z = new DenseVector(residuals.Count);

            Vector temp = new DenseVector(residuals.Count);
            
            // for (k = 0, 1, .... )
            int iterationNumber = 0;
            while (ShouldContinue(iterationNumber, xTemp, input, residuals))
            {
                // p_k = r_k + beta_(k-1) * (p_(k-1) - u_(k-1))
                p.Subtract(u, temp);
                residuals.AddScaledVector(beta, temp, p);

                // Solve M b_k = p_k
                m_Preconditioner.Approximate(p, temp);

                // s_k = A b_k
                matrix.Multiply(temp, s);

                // alpha_k = (r*_0 * r_k) / (r*_0 * s_k)
                alpha = rDash.DotProduct(residuals) / rDash.DotProduct(s);

                // y_k = t_(k-1) - r_k - alpha_k * w_(k-1) + alpha_k s_k
                s.Subtract(w, temp);
                t.Subtract(residuals, y);
                y.AddScaledVector(alpha, temp);

                // Store the old value of t in t0
				t.CopyTo(t0);

                // t_k = r_k - alpha_k s_k
                residuals.AddScaledVector(-alpha, s, t);

                // Solve M d_k = t_k
                m_Preconditioner.Approximate(t, temp);

                // c_k = A d_k
                matrix.Multiply(temp, c);
                
                double cDot = c.DotProduct();
                // cDot can only be zero if c is a zero vector
                // We'll set cDot to 1 if it is zero to prevent NaN's
                // Note that the calculation should continue fine because
                // c.DotProduct(t) will be zero and so will c.DotProduct(y)
                if (Precision.EqualsWithTolerance(cDot, 0, Constants.UlpsForComparison))
                {
                    cDot = 1.0;
                }

                // Even if we don't want to do any BiCGStab steps we'll still have
                // to do at least one at the start to initialize the
                // system, but we'll only have to take special measures
                // if we don't do any so ...
                double ctDot = c.DotProduct(t);
                if (((m_NumberOfBiCgStabSteps == 0) && (iterationNumber == 0)) || 
                    (ShouldRunBiCGStabSteps(iterationNumber)))
                {
                    // sigma_k = (c_k * t_k) / (c_k * c_k)
                    sigma = ctDot / cDot;

                    // eta_k = 0
                    eta = 0;
                }
                else
                {
                    double yDot = y.DotProduct();
                    // yDot can only be zero if y is a zero vector
                    // We'll set yDot to 1 if it is zero to prevent NaN's
                    // Note that the calculation should continue fine because
                    // y.DotProduct(t) will be zero and so will c.DotProduct(y)
                    if (Precision.EqualsWithTolerance(yDot, 0, Constants.UlpsForComparison))
                    {
                        yDot = 1.0;
                    }
                    
                    double ytDot = y.DotProduct(t);
                    double cyDot = c.DotProduct(y);

					double denom = cDot * yDot - cyDot * cyDot;
					
                    // sigma_k = ((y_k * y_k)(c_k * t_k) - (y_k * t_k)(c_k * y_k)) / ((c_k * c_k)(y_k * y_k) - (y_k * c_k)(c_k * y_k))
                    sigma = (yDot * ctDot - ytDot * cyDot ) / (denom);

                    // eta_k = ((c_k * c_k)(y_k * t_k) - (y_k * c_k)(c_k * t_k)) / ((c_k * c_k)(y_k * y_k) - (y_k * c_k)(c_k * y_k))
                    eta = (cDot * ytDot - cyDot * ctDot) / (denom);
                }

                // u_k = sigma_k s_k + eta_k (t_(k-1) - r_k + beta_(k-1) u_(k-1))
                t0.AddScaledVector(beta, u, temp);
                temp.Subtract(residuals);
                temp.Multiply(eta);

                temp.AddScaledVector(sigma, s, u);

                // z_k = sigma_k r_k +_ eta_k z_(k-1) - alpha_k u_k
                z.Multiply(eta);
                z.AddScaledVector(-alpha, u);
                z.AddScaledVector(sigma, residuals);
                
                // x_(k+1) = x_k + alpha_k p_k + z_k
                xTemp.AddScaledVector(alpha, p);
                xTemp.Add(z);
                
                // r_(k+1) = t_k - eta_k y_k - sigma_k c_k
                // Copy the old residuals to a temp vector because we'll
                // need those in the next step
                residuals.CopyTo(t0);

                t.AddScaledVector(-eta, y, residuals);
                residuals.AddScaledVector(-sigma, c);

                // beta_k = alpha_k / sigma_k * (r*_0 * r_(k+1)) / (r*_0 * r_k)
                // But first we check if there is a possible NaN. If so just reset beta to zero.
                beta = (!Precision.EqualsWithTolerance(sigma, 0, Constants.UlpsForComparison)) ? 
                    alpha / sigma * (rDash.DotProduct(residuals)) / (rDash.DotProduct(t0)) : 0;
               
                // w_k = c_k + beta_k s_k
                c.AddScaledVector(beta, s, w);

                // Get the real value
                m_Preconditioner.Approximate(xTemp, result);

                // Now check for convergence
                if (!ShouldContinue(iterationNumber, result, input, residuals))
                {
                    // Recalculate the residuals and go round again. This is done to ensure that
                    // we have the proper residuals.
                    CalculateTrueResidual(matrix, residuals, result, input);
                }

                // Next iteration.
                iterationNumber++;
            }
        }

        // Calculates the true residual of the matrix equation Ax = b according to:
        // residual = b - Ax
        private void CalculateTrueResidual(Matrix matrix, Vector residual, Vector x, Vector b)
        {
            // -Ax = residual
            matrix.Multiply(x, residual);
            residual.Negate();

            // residual + b
            residual.Add(b);
        }

        private bool ShouldContinue(int iterationNumber, Vector result, Vector source, Vector residuals)
        {
            if (m_HasBeenStopped)
            {
                m_Iterator.IterationCancelled();
                return true;
            }

            m_Iterator.DetermineStatus(iterationNumber, result, source, residuals);
            ICalculationStatus status = m_Iterator.Status;

            // We stop if either:
            // - the user has stopped the calculation
            // - the calculation needs to be stopped from a numerical point of view (divergence, convergence etc.)
            return (!status.TerminatesCalculation) && (!m_HasBeenStopped);
        }
        
        private bool ShouldRunBiCGStabSteps(int iterationNumber)
        {
            // Run the first steps as BiCGStab
            // The number of steps past a whole iteration set
            int difference = iterationNumber % (m_NumberOfBiCgStabSteps + m_NumberOfGPBiCGSteps);

            // Do steps with BiCGStab if:
            // - The difference is zero or more (i.e. we have done zero or more complete cycles)
            // - The difference is less than the number of BiCGStab steps that should be taken
            return ((difference >= 0) && (difference < m_NumberOfBiCgStabSteps));
        }

        /// <summary>
        /// Solves the matrix equation AX = B, where A is the coefficient matrix, B is the
        /// solution matrix and X is the unknown matrix.
        /// </summary>
        /// <param name="matrix">The coefficient matrix, <c>A</c>.</param>
        /// <param name="input">The solution matrix, <c>B</c>.</param>
        /// <returns>The result matrix, <c>X</c>.</returns>
        public Matrix Solve(Matrix matrix, Matrix input)
        {
            if (matrix == null)
            {
                throw new ArgumentNullException("matrix");
            }
            if (input == null)
            {
                throw new ArgumentNullException("input");
            }
            Matrix result = matrix.CreateMatrix(input.Rows, input.Columns);
            Solve(matrix, input, result);
            return result;
        }

        /// <summary>
        /// Solves the matrix equation AX = B, where A is the coefficient matrix, B is the
        /// solution matrix and X is the unknown matrix.
        /// </summary>
        /// <param name="matrix">The coefficient matrix, <c>A</c>.</param>
        /// <param name="input">The solution matrix, <c>B</c>.</param>
        /// <param name="result">The result matrix, <c>X</c></param>
        public void Solve(Matrix matrix, Matrix input, Matrix result)
        {
            if (matrix == null)
            {
                throw new ArgumentNullException("matrix");
            }
            if (input == null)
            {
                throw new ArgumentNullException("input");
            }
            if (result == null)
            {
                throw new ArgumentNullException("result");
            }

            if (matrix.Rows != input.Rows || input.Rows != result.Rows || input.Columns != result.Columns)
            {
                throw new NotConformableException();
            }

            for (int column = 0; column < input.Columns; column++)
            {
                Vector solution = Solve(matrix, input.GetColumn(column));
                foreach (KeyValuePair<int, double> element in solution.GetIndexedEnumerator())
                {
                    result.ValueAt(element.Key, column, element.Value);
                }
            }
        }
    }
}
